package com.gitee.cui.security.keystore;

import com.gitee.cui.keystore.CertHelper;
import com.gitee.cui.util.ByteUtil;
import lombok.Data;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.lang.StringUtils;
import sun.misc.BASE64Encoder;
import sun.security.provider.X509Factory;

import java.io.*;
import java.security.KeyStore;
import java.security.MessageDigest;
import java.security.PrivateKey;
import java.security.cert.Certificate;
import java.security.cert.CertificateException;
import java.security.cert.CertificateFactory;
import java.security.cert.X509Certificate;
import java.util.Base64;
import java.util.Collections;
import java.util.Date;
import java.util.List;

/**
 * 证书文件存储生成类
 * @author cuiqingqiang
 * @version 1.0
 * @date 2021/11/5 9:43
 **/
@Slf4j
@Data
public class KeyStoreHelper {

    private static final byte[] CRLF = new byte[] {'\r', '\n'};
    private static final String KEY_TYPE = "JKS";

    /**
     * 密钥存储的文件
     */
    private String keyStoreFile;

    /**
     * 获取keystore的信息所需要的密码
     */
    private String storePass;

    /**
     * 设置指定别名条目的密码，也就是私钥密码
     */
    private String keyPass;

    /**
     * 每个keystore都关联一个独一无二的别名，不区分大小写
     */
    private String alias;

    /**
     * 指定整数拥有者的信息
     * 例如："CN=名字与姓氏,OU=组织单位名称,O=组织名称,L=城市或区域名称,ST=州或省份名称,C=单位的两字母国家代码"
     */
    private String dname;

    /**
     * 证书顶层类定义
     */
    private KeyStore keyStore;

    public KeyStoreHelper (String keyStoreFile, String storePass, String keyPass, String alias, String dname) {
        this.keyStoreFile = keyStoreFile;
        this.storePass = storePass;
        this.keyPass = keyPass;
        this.alias = alias;
        this.dname = dname;
    }

    /**
     * 创建密钥和证书并保存到密钥仓库文件
     * @throws Exception 创建失败
     */
    public void createKeyEntry() throws Exception {
        // 从证书文件中读取证书内容，如果证书不存在，则创建一个空证书
        KeyStore ks = loadStore();
        CertHelper certHelper = new CertHelper(dname);
        // 生成证书
        Certificate cert = certHelper.genCert();
        cert.verify(certHelper.getKeyPair().getPublic());
        PrivateKey privateKey = certHelper.getKeyPair().getPrivate();
        char[] passArray = storePass.toCharArray();

        // 设置密钥和证书到密钥仓库
        ks.setKeyEntry(alias, privateKey, passArray, new Certificate[]{cert});

        // 生成到文件中
        try (FileOutputStream fos = new FileOutputStream(keyStoreFile)) {
            ks.store(fos, passArray);
        }
    }

    /**
     * 从文件加载keystore密钥仓库
     * @return 密钥仓库
     */
    private KeyStore loadStore() throws Exception {
        log.debug("keyStoreFile: {}", keyStoreFile);
        if (!new File(keyStoreFile).exists()) {
            createEmptyKeyStore();
        }
        KeyStore ks = KeyStore.getInstance(KEY_TYPE);
        // 从文件中读取密钥仓库
        try (FileInputStream fis = new FileInputStream(keyStoreFile)) {
            ks.load(fis, storePass.toCharArray());
        }
        return ks;
    }

    /**
     * 创建一个空的KeyStore仓库
     */
    private void createEmptyKeyStore() throws Exception {
        KeyStore ks = KeyStore.getInstance(KEY_TYPE);
        File parentFile = new File(keyStoreFile).getParentFile();
        if (!parentFile.exists()) {
            parentFile.mkdirs();
        }
        ks.load(null, storePass.toCharArray());
        try (FileOutputStream fos = new FileOutputStream(keyStoreFile)) {
            ks.store(fos, storePass.toCharArray());
        }
    }

    /**
     * 打印密钥仓库文件中的所有信息
     */
    public void printEntries() throws Exception {
        if (keyStore == null) {
            keyStore = loadStore();
        }
        List<String> aliases = Collections.list(keyStore.aliases());
        aliases.sort(String::compareTo);
        for (String alias : aliases) {
            doPrintEntry(alias);
        }
    }

    /**
     * 打印单个keystore信息
     * @param alias keystore别名
     */
    private void doPrintEntry(String alias) throws Exception {
        log.info("{} 别名的证书信息如下：", alias);
        Certificate[] chain = keyStore.getCertificateChain(alias);
        if (chain == null) {
            if (keyStore.entryInstanceOf(alias, KeyStore.TrustedCertificateEntry.class)) {
                Certificate certificate = keyStore.getCertificate(alias);
                printX509Cert((X509Certificate) certificate);
            }
            return;
        }

        for (Certificate certificate : chain) {
            if (certificate instanceof X509Certificate) {
                printX509Cert((X509Certificate) certificate);
            } else {
                dumpCert(certificate);
            }
        }
    }

    /**
     * 将证书内容通过base64方式打印出来
     * @param certificate 证书信息
     * @throws Exception 编码错误
     */
    private void dumpCert(Certificate certificate) throws Exception {
        log.info(X509Factory.BEGIN_CERT);
        log.info(Base64.getMimeEncoder(64, CRLF).encodeToString(certificate.getEncoded()));
        log.info(X509Factory.END_CERT);
    }

    /**
     * 打印X509证书内容
     * @param cert 证书信息
     */
    private void printX509Cert(X509Certificate cert) throws Exception {
        String sigName = cert.getSigAlgName();
        log.info("Owner: {}", cert.getSubjectDN());
        log.info("Issuer: {}", cert.getIssuerDN());
        log.info("Serial number: {}", cert.getSerialNumber());
        log.info("Valid from: {}", cert.getNotBefore());
        log.info("Valid until: {}", cert.getNotAfter());
        log.info("Certificate fingerprints SHA1: ");
        log.info(getCertFingerPrint("SHA-1", cert));
        log.info("Certificate fingerprints SHA256: ");
        log.info(getCertFingerPrint("SHA-256", cert));
        log.info("Signature algorithm name: {}", sigName);
        log.info("Version: {}", cert.getVersion());
    }

    /**
     * 获取证书指纹信息
     * @param s 加密方式
     * @param cert 证书信息
     * @return 对应加密方式的指纹信息
     */
    private String getCertFingerPrint(String s, X509Certificate cert) throws Exception {
        byte[] bytes = cert.getEncoded();
        MessageDigest digest = MessageDigest.getInstance(s);
        byte[] result = digest.digest(bytes);
        return ByteUtil.byteToHex(result);
    }

    /**
     * 导出证书
     * @param outDir 到处的目标目录
     * @return 是否导出成功
     */
    public boolean exportCert(String outDir) throws Exception {
        assert (StringUtils.isNotEmpty(alias));
        assert (StringUtils.isNotEmpty(keyPass));

        KeyStore ks = loadStore();
        KeyStore.PasswordProtection protection = new KeyStore.PasswordProtection(keyPass.toCharArray());
        if (ks.isKeyEntry(alias)) {
            KeyStore.PrivateKeyEntry entry = (KeyStore.PrivateKeyEntry) ks.getEntry(alias, protection);
            X509Certificate certificate = (X509Certificate) entry.getCertificate();

            // 判断证书是否过期
            if (new Date().after(certificate.getNotAfter())) {
                return false;
            }
            // 导出到文件
            String certPath = outDir + File.separator + alias + ".cer";
            try (FileWriter fw = new FileWriter(new File(certPath))) {
                String encode = new BASE64Encoder().encode(certificate.getEncoded());
                String result = "-----BEGIN CERTIFICATE-----\r\n" + encode + "\r\n-----END CERTIFICATE-----\r\n";
                fw.write(result);
                fw.flush();
                return true;
            }
        }
        return false;
    }

    /**
     * 导入数字证书到信任仓库
     * @param importAlias 导入数字证书别名
     * @param certPath 证书路径
     * @return 导入是否成功
     */
    public boolean importCert(String importAlias, String certPath) throws Exception {
        if (keyStore == null) {
            keyStore = loadStore();
        }
        FileInputStream fis = null;
        if (StringUtils.isNotEmpty(certPath)) {
            File file;
            fis = new FileInputStream(certPath);
        }
        return addTrustedCert(importAlias, fis);
    }

    /**
     * 将证书添加到信任仓库
     * @param importAlias 导入别名
     * @param fis 证书输入流
     * @return 是否导入成功
     */
    private boolean addTrustedCert(String importAlias, FileInputStream fis) throws Exception {
        if (StringUtils.isBlank(importAlias)) {
            throw new Exception("Must specify alias");
        }

        if (keyStore.containsAlias(importAlias)) {
            throw new Exception("别名已经存在");
        }
        // 从输入流读取证书内容
        X509Certificate cert = null;
        try {
            cert = (X509Certificate) generateCertificate(fis);
        } catch (Exception e) {
            throw new Exception("证书读取失败");
        }

        // 根据别名进行设置
        keyStore.setCertificateEntry(importAlias, cert);
        // 保存到文件
        char[] array = storePass.toCharArray();
        try (FileOutputStream outputStream = new FileOutputStream(keyStoreFile)) {
            keyStore.store(outputStream, array);
        }
        return true;
    }

    /**
     * 创建证书
     * @param fis 输入流
     * @return 证书内容
     */
    private Certificate generateCertificate(FileInputStream fis) throws CertificateException {
        byte[] data = readInputStream(fis);
        return CertificateFactory.getInstance("X.509").generateCertificate(new ByteArrayInputStream(data));
    }

    /**
     * 将文件流转换为字节数组
     * @param fis 文件输入流
     * @return 字节数组
     */
    private byte[] readInputStream(FileInputStream fis) {
        ByteArrayOutputStream out = new ByteArrayOutputStream();
        byte[] buffer = new byte[1024];
        int length;
        try {
            while ((length = fis.read(buffer)) != -1) {
                out.write(buffer, 0, length);
            }
        } catch (IOException e) {
            e.printStackTrace();
        }
        return out.toByteArray();
    }
}
